# import torch
# from torchvision import models, transforms
# from torch.autograd import Variable
# from PIL import Image
# import matplotlib.pyplot as plt
# import numpy as np
# import joblib
# import os
#
# # 加载模型和标签编码器
# vgg16 = models.vgg16(pretrained=False)
# vgg16.classifier[6] = torch.nn.Linear(vgg16.classifier[6].in_features, 3)
# vgg16.load_state_dict(torch.load('vgg16_best_model.pth'))
# vgg16 = vgg16.to('cuda' if torch.cuda.is_available() else 'cpu')
# vgg16.eval()
# # label_encoder = joblib.load('label_encoder.pkl')
# label_encoder = ['最优级-1级', '中间级-2级', '差等级-3级']
#
# # 图像预处理
# transform = transforms.Compose([
#     transforms.Resize((224, 224)),
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# ])
#
# # 提取中间层特征
# class FeatureExtractor:
#     def __init__(self, model, layers):
#         self.model = model
#         self.layers = layers
#         self.outputs = {}
#         self._register_hooks()
#
#     def _register_hooks(self):
#         for layer_name in self.layers:
#             layer = dict([*self.model.named_modules()])[layer_name]
#             layer.register_forward_hook(self._hook_fn(layer_name))
#
#     def _hook_fn(self, layer_name):
#         def hook(module, input, output):
#             self.outputs[layer_name] = output
#         return hook
#
#     def __call__(self, x):
#         self.model(x)
#         return self.outputs
#
# # 预测函数并可视化特征
# def predict_and_visualize(image_path, save_dir='features'):
#     if not os.path.exists(save_dir):
#         os.makedirs(save_dir)
#
#     image = Image.open(image_path).convert('RGB')
#     image_tensor = transform(image).unsqueeze(0)
#     image_tensor = image_tensor.to('cuda' if torch.cuda.is_available() else 'cpu')
#
#     layers = [
#         'features.0', 'features.2', 'features.5', 'features.7',
#         'features.10', 'features.12', 'features.14', 'features.17',
#         'features.19', 'features.21', 'features.24', 'features.26',
#         'features.28'
#     ]
#     feature_extractor = FeatureExtractor(vgg16, layers)
#
#     features = feature_extractor(image_tensor)
#
#     # 预测类别
#     outputs = vgg16(image_tensor)
#     _, predicted = torch.max(outputs, 1)
#     predicted_class = label_encoder[predicted.item()]
#
#     # 可视化特征并保存
#     for layer_name, feature in features.items():
#         feature = feature.cpu().data.numpy()
#         fig, axarr = plt.subplots(min(8, feature.shape[1]) // 4, 4, figsize=(15, 15))
#         fig.suptitle(f'Layer {layer_name}')
#         for idx in range(min(8, feature.shape[1])):  # 仅显示前8个特征图
#             ax = axarr[idx // 4, idx % 4]
#             ax.imshow(feature[0, idx], cmap='gray')
#             ax.axis('off')
#         plt.savefig(os.path.join(save_dir, f'{layer_name}.png'))
#         plt.close(fig)
#
#     return predicted_class
#
# # 示例用法
# image_path = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_1\1.tif'  # 替换为实际的图片路径
# predicted_class = predict_and_visualize(image_path)
# print("Predicted Class:", predicted_class)


import torch
from torchvision import models, transforms
from torch.autograd import Variable
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import joblib
import os

# 加载模型和标签编码器
vgg16 = models.vgg16(pretrained=False)
vgg16.classifier[6] = torch.nn.Linear(vgg16.classifier[6].in_features, 3)
vgg16.load_state_dict(torch.load('vgg16_best_model.pth'))
vgg16 = vgg16.to('cuda' if torch.cuda.is_available() else 'cpu')
vgg16.eval()
# label_encoder = joblib.load('label_encoder.pkl')
label_encoder = ['最优级-1级', '中间级-2级', '差等级-3级']


# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 提取中间层特征
class FeatureExtractor:
    def __init__(self, model, layers):
        self.model = model
        self.layers = layers
        self.outputs = {}
        self._register_hooks()

    def _register_hooks(self):
        for layer_name in self.layers:
            layer = dict([*self.model.named_modules()])[layer_name]
            layer.register_forward_hook(self._hook_fn(layer_name))

    def _hook_fn(self, layer_name):
        def hook(module, input, output):
            self.outputs[layer_name] = output
        return hook

    def __call__(self, x):
        self.model(x)
        return self.outputs

# 预测函数并可视化特征
def predict_and_visualize(image_path, save_dir='features2'):
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    image = Image.open(image_path).convert('RGB')
    image_tensor = transform(image).unsqueeze(0)
    image_tensor = image_tensor.to('cuda' if torch.cuda.is_available() else 'cpu')

    layers = [
        'features.0', 'features.2', 'features.5', 'features.7',
        'features.10', 'features.12', 'features.14', 'features.17',
        'features.19', 'features.21', 'features.24', 'features.26',
        'features.28'
    ]
    feature_extractor = FeatureExtractor(vgg16, layers)

    features = feature_extractor(image_tensor)

    # 预测类别
    outputs = vgg16(image_tensor)
    _, predicted = torch.max(outputs, 1)
    predicted_class = label_encoder[predicted.item()]

    # 可视化特征并保存
    for layer_name, feature in features.items():
        feature = feature.cpu().data.numpy()
        num_features = min(8, feature.shape[1])  # 仅显示前8个特征图
        for idx in range(num_features):
            plt.figure(figsize=(2, 2))
            plt.imshow(feature[0, idx], cmap='Blues')
            plt.axis('off')
            plt.savefig(os.path.join(save_dir, f'{layer_name}_{idx}.png'))
            plt.close()

    return predicted_class

# 示例用法
image_path = r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_3\1.tif'  # 替换为实际的图片路径
predicted_class = predict_and_visualize(image_path)
print("Predicted Class:", predicted_class)
